package com.jtzc.aikf.service.impl;

import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.jtzc.aikf.entity.DocumentDO;
import com.jtzc.aikf.service.DocumentService;
import com.jtzc.aikf.service.KnowledgeService;
import com.jtzc.aikf.service.VectorService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.UUID;

/**
 * @author wu chuang
 * @description
 */
@Service
public class KnowledgeServiceImpl implements KnowledgeService {

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private EmbeddingStore embeddingStore;

    @Autowired
    private DocumentService documentService;

    @Override
    public Page<DocumentDO> page(Page<DocumentDO> objectPage) {
        Page<DocumentDO> page = documentService.page(objectPage);
        return page;
    }

    @Override
    public void upload(MultipartFile file) throws Exception {
        try {
            // 获取上传目录路径
            String uploadDir = System.getProperty("user.dir") + "/uploads/";
            File dir = new File(uploadDir);
            if (!dir.exists()) dir.mkdirs();

            // 生成唯一文件名
            String filename = UUID.randomUUID() + "_" + file.getOriginalFilename();
            // 保存文件
            file.transferTo(new File(uploadDir + filename));
            // 保存文档
            DocumentDO documentDO = new DocumentDO();
            documentDO.setFileName(filename);
            documentDO.setFilePath("/uploads/"+filename);
            documentDO.setFileType(file.getContentType());
            documentDO.setKnowledgeBaseId(1);
            documentDO.setCreateBy("张三");
            documentDO.setUpdateBy("李四");
            documentDO.setCreateTime(new Date());
            documentDO.setUpdateTime(new Date());
            documentService.save(documentDO);
            // 向量存储
            Document document = FileSystemDocumentLoader.loadDocument(uploadDir + filename);
            List<Document> documents = Arrays.asList(document);

            //文本向量化并存入向量数据库：将每个片段进行向量化，得到一个嵌入向量
            EmbeddingStoreIngestor
                    .builder()
                    .embeddingStore(embeddingStore)
                    .embeddingModel(embeddingModel)
                    .build()
                    .ingest(documents);
        } catch (Exception e) {
            throw e;
        }
    }

    @Autowired
    private VectorService vectorService;
    @Override
    public void relearn(Integer docId) {
        DocumentDO document = documentService.getById(docId);
        vectorService.learn(document);
    }
}
